PySpark Job Examples

To execute PySpark jobs on Amazon EMR, you must construct your script to correctly integrate with S3 (via EMRFS) and configure your cluster's resources. This guide shows a production-grade PySpark ETL script and the matching shell submit script.

1. PySpark ETL Script (emr_pyspark_etl.py)

This script reads a raw CSV dataset from an S3 bucket, performs transformations (aggregations and filtering), and writes the output back to S3 in partitioned Parquet format.

import sys
import argparse
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum as spark_sum, current_timestamp

def run_etl(input_path, output_path):
    # Initialize Spark Session (Configured automatically by YARN in EMR)
    spark = SparkSession.builder \
        .appName("EMR-Production-PySpark-ETL") \
        .getOrCreate()

    print(f"Starting Spark job. Input path: {input_path}")
    print(f"Output will be saved to: {output_path}")

    # 1. Read CSV raw data from S3 using EMRFS
    # Header is True, inferring schema automatically
    raw_df = spark.read.option("header", "true") \
                       .option("inferSchema", "true") \
                       .csv(input_path)

    print("Raw Data Schema:")
    raw_df.printSchema()

    # 2. Transform: Filter, Aggregate, and Add metadata
    # Example dataset schema: country, industry, revenue, year
    transformed_df = raw_df.filter(col("revenue").isNotNull()) \
        .groupBy("country", "industry", "year") \
        .agg(
            spark_sum("revenue").alias("total_revenue")
        ) \
        .withColumn("processed_at", current_timestamp())

    # 3. Write data back to S3 in Parquet format, partitioned by year
    print("Writing processed Parquet files back to S3...")
    transformed_df.write \
        .mode("overwrite") \
        .partitionBy("year") \
        .parquet(output_path)

    print("Job completed successfully!")
    spark.stop()

if __name__ == "__main__":
    # Parsing command-line arguments passed by EMR Step / Spark Submit
    parser = argparse.ArgumentParser(description="EMR PySpark ETL Script")
    parser.add_argument("--input", required=True, help="S3 URI for input CSV data")
    parser.add_argument("--output", required=True, help="S3 URI for output Parquet data")

    args = parser.parse_args()

    run_etl(args.input, args.output)

2. EMR Spark Submit Command (submit_job.sh)

You can submit your script to EMR from inside the Primary node or using EMR Steps. The standard deployment configuration is to execute the job in cluster mode so the Driver runs on a worker node rather than overloading the Primary node.

#!/bin/bash

# Configuration
SCRIPT_S3_URI="s3://my-spark-jobs-bucket/scripts/emr_pyspark_etl.py"
INPUT_S3_URI="s3://my-spark-jobs-bucket/raw-data/orders/"
OUTPUT_S3_URI="s3://my-spark-jobs-bucket/processed-data/orders_summary/"

# Spark-Submit execution matching EMR resources
spark-submit \
    --master yarn \
    --deploy-mode cluster \
    --num-executors 10 \
    --executor-cores 4 \
    --executor-memory 16g \
    --driver-memory 8g \
    --conf spark.dynamicAllocation.enabled=true \
    --conf spark.serializer=org.apache.spark.serializer.KryoSerializer \
    $SCRIPT_S3_URI \
    --input $INPUT_S3_URI \
    --output $OUTPUT_S3_URI

3. Best Practices for PySpark on EMR

1. Dynamic Resource Allocation: Instead of hardcoding --num-executors, set --conf spark.dynamicAllocation.enabled=true so YARN can provision and de-provision executor resources automatically depending on task queue length.
2. Writing to S3 Direct (Multipart Uploader): S3 write performance can be slow if Spark has to write temporary files and rename them. To optimize this, EMR includes the EMRFS S3 Optimized Committer automatically when you write to Parquet. Avoid disabling it as it significantly speeds up job write operations.
3. Partitioning: Always partition your final output datasets by logical dimensions (such as year, month, or category) when writing to S3. This enables downstream engines (Athena, Spark, Redshift) to prune partitions during query execution, reducing S3 scan volume and accelerating query times.